// Copyright (c) 2022 Presto Labs Pte. Ltd.
// Author: fengyang

#include <memory>

#include <boost/asio/io_service.hpp>

#include "coin2/strategy/strategy.h"
#include "coin2/exchange/base/feed/timer/deadline_timer.h"

#include <gtest/gtest.h>

using namespace coin2::exchange::base::feed::timer;

using coin2::exchange::base::feed::IFeedSubscriber;
using coin2::exchange::base::feed::FeedUpdate;

class PsudoStrategy : public ::coin2::strategy::IStrategy {
 public:
  void Init(::coin2::app::Driver* driver) {
    timer1_ = std::make_unique<DeadlineTimer>();
    timer2_ = std::make_unique<DeadlineTimer>();
    timer3_ = std::make_unique<DeadlineTimer>();
    on_heartbeat_callback_ = []() {}; // empty
  }

  void onHeartbeatFeed(const FeedTimer& timer) override {
    last_feed_timer_ = timer;
    timer_call_count++;
    on_heartbeat_callback_();
  }

  DeadlineTimer* GetTimer1() {
    return timer1_.get();
  }

  DeadlineTimer* GetTimer2() {
    return timer2_.get();
  }

  DeadlineTimer* GetTimer3() {
    return timer3_.get();
  }

  FeedTimer* GetFeedTimer() {
    return &last_feed_timer_;
  }

  int64_t GetTimerCount() {
    return timer_call_count;
  }

  void SetCallback(std::function<void(void)> callback) {
    on_heartbeat_callback_ = callback;
  }

  void ResetCallback() {
    on_heartbeat_callback_ = []() {};
  }

 private:
  std::unique_ptr<DeadlineTimer> timer1_;
  std::unique_ptr<DeadlineTimer> timer2_;
  std::unique_ptr<DeadlineTimer> timer3_;
  FeedTimer last_feed_timer_;
  int64_t timer_call_count = 0;
  std::function<void(void)> on_heartbeat_callback_;
};


// change visibility for testing
class PsudoArchiveTimerManager: public ArchiveTimerManager {
 public:
  virtual ~PsudoArchiveTimerManager() {}
  PsudoArchiveTimerManager() {}
  static std::shared_ptr<PsudoArchiveTimerManager> CreateArchiveManager() {
    CHECK(!instance_);
    std::shared_ptr<PsudoArchiveTimerManager> ptr(new PsudoArchiveTimerManager());
    instance_ = ptr;
    return ptr;
  }
  static void DeleteInstance() {
    return ArchiveTimerManager::DeleteInstance();
  }

  IFeedSubscriber* HookSubscriber(IFeedSubscriber* subscriber) {
    return ArchiveTimerManager::HookSubscriber(subscriber);
  }
};

// change visibility for testing
class PsudoLiveTimerManager: public LiveTimerManager {
 public:
  virtual ~PsudoLiveTimerManager() {}
  PsudoLiveTimerManager() {}
  static std::shared_ptr<PsudoLiveTimerManager> CreateLiveManager() {
    CHECK(!instance_);
    std::shared_ptr<PsudoLiveTimerManager> ptr(new PsudoLiveTimerManager());
    instance_ = ptr;
    return ptr;
  }
  static void DeleteInstance() {
    return LiveTimerManager::DeleteInstance();
  }

  void RegisterSubscriber(IFeedSubscriber* subscriber) {
    return LiveTimerManager::RegisterSubscriber(subscriber);
  }

  void RegisterTimerIoc(boost::asio::io_context& ioc) {
    return LiveTimerManager::RegisterTimerIoc(ioc);
  }

};

class DeadlineTimerArchiveTest : public ::testing::Test {
 protected:
  void SetUp() override {
    auto manager = PsudoArchiveTimerManager::CreateArchiveManager();
    subscriber_ = manager->HookSubscriber(&strategy_);
    strategy_.Init(nullptr);

    std::string system_name;
    MarketExchangeApi mea;
    feed_update_ = std::make_unique<FeedUpdate>(system_name, mea);
  }

  void TearDown() override {
    PsudoArchiveTimerManager::DeleteInstance();
  }

  PsudoStrategy strategy_;
  IFeedSubscriber* subscriber_;
  std::unique_ptr<FeedUpdate> feed_update_;
};

class DeadlineTimerLiveTest : public ::testing::Test {
 protected:
  void SetUp() override {
    main_work_ = std::make_unique<boost::asio::io_context::work>(main_ioc_);
    auto manager = PsudoLiveTimerManager::CreateLiveManager();
    manager->RegisterSubscriber(&strategy_);
    manager->RegisterTimerIoc(main_ioc_);
    timer_thread_ = std::thread([&]() { main_ioc_.run(); });
    strategy_.Init(nullptr);
  }

  void TearDown() override {
    main_work_.reset();
    main_ioc_.stop();
    timer_thread_.join();
    PsudoLiveTimerManager::DeleteInstance();
  }

  PsudoStrategy strategy_;
  IFeedSubscriber* subscriber_;
  boost::asio::io_context main_ioc_;
  std::unique_ptr<boost::asio::io_context::work> main_work_;
  std::thread timer_thread_;
};


TEST_F(DeadlineTimerArchiveTest, basic) {
  strategy_.GetTimer1()->ExpiresFromNow(3'000'000'000LL);
  int64_t now = 1657511745000000000;
  feed_update_->SetTimestamp(now);
  subscriber_->onBookFeed(*feed_update_);
  EXPECT_EQ(strategy_.GetFeedTimer()->TimerId(), 0);
  EXPECT_EQ(strategy_.GetFeedTimer()->Timestamp(), 0);

  strategy_.GetTimer1()->ExpiresFromNow(6'000'000'000LL);
  feed_update_->SetTimestamp(now + 5'000'000'000LL);
  subscriber_->onBookFeed(*feed_update_);
  EXPECT_EQ(strategy_.GetFeedTimer()->TimerId(), 0);
  EXPECT_EQ(strategy_.GetFeedTimer()->Timestamp(), 0);

  feed_update_->SetTimestamp(now + 9'000'000'000LL);
  subscriber_->onBookFeed(*feed_update_);
  EXPECT_EQ(strategy_.GetFeedTimer()->Timestamp(), now + 6'000'000'000LL);

  strategy_.GetTimer1()->ExpiresAt(now + 15'000'000'000LL);
  feed_update_->SetTimestamp(now + 15'000'000'000LL);
  subscriber_->onBookFeed(*feed_update_);
  EXPECT_EQ(strategy_.GetFeedTimer()->Timestamp(), now + 15'000'000'000LL);

  strategy_.GetTimer1()->ExpiresEvery(3'000'000'000LL);
  feed_update_->SetTimestamp(now + 18'000'000'000LL);
  subscriber_->onBookFeed(*feed_update_);
  EXPECT_EQ(strategy_.GetFeedTimer()->Timestamp(), now + 18'000'000'000LL);
  feed_update_->SetTimestamp(now + 20'000'000'000LL);
  subscriber_->onBookFeed(*feed_update_);
  EXPECT_EQ(strategy_.GetFeedTimer()->Timestamp(), now + 18'000'000'000LL);

  feed_update_->SetTimestamp(now + 21'000'000'000LL);
  subscriber_->onBookFeed(*feed_update_);
  EXPECT_EQ(strategy_.GetFeedTimer()->TimerId(), 1);
  EXPECT_EQ(strategy_.GetTimerCount(), 4);
}


// test timer expiresXXX inside onHeartbeatFeed
TEST_F(DeadlineTimerArchiveTest, inside) {
  int64_t now = 1657511745000000000;
  strategy_.GetTimer1()->ExpiresFromNow(1'000'000'000LL);

  feed_update_->SetTimestamp(now);
  subscriber_->onBookFeed(*feed_update_);

  // not trigger
  EXPECT_EQ(strategy_.GetFeedTimer()->TimerId(), 0);
  EXPECT_EQ(strategy_.GetFeedTimer()->Timestamp(), 0);

  strategy_.SetCallback(std::bind(
    [](DeadlineTimer* timer, int64_t timestamp){ timer->ExpiresFromNow(timestamp); },
    strategy_.GetTimer1(),
    2'000'000'000LL
  ));

  feed_update_->SetTimestamp(now + 3'000'000'000LL);
  subscriber_->onBookFeed(*feed_update_);
  EXPECT_EQ(strategy_.GetTimerCount(), 2);

  strategy_.SetCallback(std::bind(
    [](DeadlineTimer* timer, int64_t timestamp){
        static bool running = false;
        if (running) return;
        running = true;
        timer->ExpiresAt(timestamp); },
    strategy_.GetTimer1(),
    now + 7'000'000'000LL
  ));
  feed_update_->SetTimestamp(now + 8'000'000'000LL);
  subscriber_->onBookFeed(*feed_update_);
  EXPECT_EQ(strategy_.GetTimerCount(), 4);

  feed_update_->SetTimestamp(now + 10'000'000'000LL);
  subscriber_->onBookFeed(*feed_update_);
  EXPECT_EQ(strategy_.GetTimerCount(), 4);
}


TEST_F(DeadlineTimerArchiveTest, multiple) {
  strategy_.GetTimer1()->ExpiresEvery(3'000'000'000LL);
  strategy_.GetTimer2()->ExpiresEvery(5'000'000'000LL);

  int64_t now = 1657511745000000000;
  feed_update_->SetTimestamp(now);
  subscriber_->onBookFeed(*feed_update_);

  // not trigger
  EXPECT_EQ(strategy_.GetFeedTimer()->TimerId(), 0);
  EXPECT_EQ(strategy_.GetFeedTimer()->Timestamp(), 0);

  feed_update_->SetTimestamp(now + 5'000'000'000LL);
  subscriber_->onBookFeed(*feed_update_);

  strategy_.GetTimer3()->ExpiresFromNow(2'000'000'000LL);

  feed_update_->SetTimestamp(now + 10'000'000'000LL);
  subscriber_->onBookFeed(*feed_update_);
  EXPECT_EQ(strategy_.GetTimerCount(), 6);

  strategy_.GetTimer3()->ExpiresAt(now + 13'000'000'000LL);

  feed_update_->SetTimestamp(now + 15'000'000'000LL);
  subscriber_->onBookFeed(*feed_update_);
  EXPECT_EQ(strategy_.GetTimerCount(), 10);
}


TEST_F(DeadlineTimerLiveTest, basic) {
  using namespace std::chrono_literals;

  auto now = GetCurrentTimestamp();
  strategy_.GetTimer1()->ExpiresFromNow(30'000'000LL);
  std::this_thread::sleep_for(10ms);
  EXPECT_EQ(strategy_.GetFeedTimer()->TimerId(), 0);
  EXPECT_EQ(strategy_.GetFeedTimer()->Timestamp(), 0);

  now = GetCurrentTimestamp();
  strategy_.GetTimer1()->ExpiresFromNow(60'000'000LL);
  std::this_thread::sleep_for(50ms);
  EXPECT_EQ(strategy_.GetFeedTimer()->TimerId(), 0);
  EXPECT_EQ(strategy_.GetFeedTimer()->Timestamp(), 0);

  std::this_thread::sleep_for(20ms);
  EXPECT_LT(strategy_.GetFeedTimer()->Timestamp() - now - 60'000'000LL, 100'000LL);

  now = GetCurrentTimestamp();
  strategy_.GetTimer1()->ExpiresAt(now + 10'000'000LL);
  std::this_thread::sleep_for(10ms);
  EXPECT_LT(strategy_.GetFeedTimer()->Timestamp() - now - 10'000'000LL, 100'000LL);

  now = GetCurrentTimestamp();
  strategy_.GetTimer1()->ExpiresEvery(30'000'000LL);
  std::this_thread::sleep_for(30ms);
  EXPECT_LT(strategy_.GetFeedTimer()->Timestamp() - now - 30'000'000LL,  100'000LL);
  std::this_thread::sleep_for(30ms);
  EXPECT_LT(strategy_.GetFeedTimer()->Timestamp() - now - 2 * 30'000'000LL, 100'000LL);

  EXPECT_EQ(strategy_.GetFeedTimer()->TimerId(), 1);
  EXPECT_EQ(strategy_.GetTimerCount(), 4);
}

// test timer expiresXXX inside onHeartbeatFeed
TEST_F(DeadlineTimerLiveTest, inside) {
  using namespace std::chrono_literals;

  auto now = GetCurrentTimestamp();
  strategy_.GetTimer1()->ExpiresFromNow(10'000'000LL);
  std::this_thread::sleep_for(5ms);
  // not trigger
  EXPECT_EQ(strategy_.GetFeedTimer()->TimerId(), 0);
  EXPECT_EQ(strategy_.GetFeedTimer()->Timestamp(), 0);
  strategy_.SetCallback(std::bind(
    [](DeadlineTimer* timer, int64_t timestamp){ timer->ExpiresFromNow(timestamp); },
    strategy_.GetTimer1(),
    20'000'000LL
  ));
  std::this_thread::sleep_for(25ms);
  EXPECT_EQ(strategy_.GetTimerCount(), 2);

  strategy_.SetCallback(std::bind(
    [](DeadlineTimer* timer, int64_t timestamp){
        static bool running = false;
        if (running) return;
        running = true;
        timer->ExpiresAt(timestamp); },
    strategy_.GetTimer1(),
    now + 70'000'000LL
  ));

  std::this_thread::sleep_for(25ms);
  EXPECT_EQ(strategy_.GetTimerCount(), 3);

  std::this_thread::sleep_for(40ms);
  EXPECT_EQ(strategy_.GetTimerCount(), 4);
}

TEST_F(DeadlineTimerLiveTest, multiple) {
  using namespace std::chrono_literals;

  int64_t now = GetCurrentTimestamp();
  strategy_.GetTimer1()->ExpiresEvery(30'000'000LL);
  strategy_.GetTimer2()->ExpiresEvery(50'000'000LL);

  std::this_thread::sleep_for(10ms);
  // not trigger
  EXPECT_EQ(strategy_.GetFeedTimer()->TimerId(), 0);
  EXPECT_EQ(strategy_.GetFeedTimer()->Timestamp(), 0);

  std::this_thread::sleep_for(40ms);
  strategy_.GetTimer3()->ExpiresFromNow(20'000'000LL);

  std::this_thread::sleep_for(50ms);
  EXPECT_EQ(strategy_.GetTimerCount(), 6);

  strategy_.GetTimer3()->ExpiresAt(now + 130'000'000LL);
  std::this_thread::sleep_for(50ms);
  EXPECT_EQ(strategy_.GetTimerCount(), 10);
}
